import pandas as pd
import numpy as np
import math
import matplotlib.pyplot as plt
from scipy.spatial.kdtree import KDTree
from sklearn.cluster import DBSCAN

def readDataSet(filename, div=2):
    data = np.loadtxt(filename)
    return data

def readLabel(filename):
    data = pd.read_csv(filename, header=None)
    matrix = np.array(data)
    return matrix

def dbscan(data):
    dbsc = DBSCAN().fit(data)
    result = dbsc.labels_
    result = np.array(result)
    return result

from sklearn.decomposition import PCA

def plotResult(matrix, result, filename):
    pca = PCA(n_components=2)  # n_components设置降维后的特征数
    pca.fit(matrix)  # 拟合
    location = pca.transform(matrix)  # 获取新矩阵

    length = len(result)
    markers = ['.', '*', '+', 'x', '^']
    colors = ['maroon', 'red', 'peru', 'gold', 'olive', 'yellowgreen', 'lawngreen', 'springgreen']
    colors = colors + ['turquoise', 'teal', 'deepskyblue', 'dodgerblue', 'royalblue', 'navy']
    colors = colors + ['slategrey', 'orchid', 'm', 'deeppink', 'crimson']
    plt.figure(figsize=(10, 10))

    for i in range(0, length):
        index = int(result[i])
        if index == -1:
            plt.plot(location[i][0], location[i][1], color=(0,0,0),  marker='.')
        else:
            plt.plot(location[i][0], location[i][1], color=colors[index%19], marker=markers[index%5])
    plt.xlabel('Attribute 1'), plt.ylabel('Attribute 2')

    plt.savefig(filename)
    plt.show()



from getRI import rand_index
from getNMI import printNMI

filename = 'data.txt'
output_file = 'result.txt'
output_img = 'result.png'
data = readDataSet(filename)
result = dbscan(data)
plotResult(data, result, output_img)
np.savetxt(output_file, result)

nmi = printNMI("label.txt", output_file)
ri = rand_index("label.txt", output_file)

np.savetxt("dt={:.2f}_nmi={:.2f}_ri={:.2f}".format(dt, nmi, ri)+output_file, result)